#include "PrincipalComponentAnalysisMotionSegmenterWidget.h"
#include "ui_PrincipalComponentAnalysisMotionSegmenterWidget.h"

#include <MMM/Motion/Plugin/KinematicPlugin/KinematicSensor.h>
#include <MMM/Motion/Plugin/MoCapMarkerPlugin/MoCapMarkerSensor.h>
#include <MMMSimoxTools/MMMSimoxTools.h>
#include <VirtualRobot/Robot.h>
#include <QTreeWidget>
#include <QMessageBox>

PrincipalComponentAnalysisMotionSegmenterWidget::PrincipalComponentAnalysisMotionSegmenterWidget(MMM::MotionPtr motion, MMM::SegmentationType segType, QWidget* parent) :
    QWidget(parent),
    ui(new Ui::PrincipalComponentAnalysisMotionSegmenterWidget),
    segType(segType),
    motion(motion)
{
    ui->setupUi(this);

    if (motion) updateTreeWidget();
}

PrincipalComponentAnalysisMotionSegmenterWidget::~PrincipalComponentAnalysisMotionSegmenterWidget()
{
    delete ui;
}

MMM::PrincipalComponentAnalysisMotionSegmenterConfigurationPtr PrincipalComponentAnalysisMotionSegmenterWidget::getConfiguration() {
    float alpha = ui->alphaSpinBox->value();
    int reducedDimensions = ui->reducedDimensionsSpinBox->value();
    MMM::KeyframeType kType;
    if (ui->firstKeyframeRadioButton->isChecked()) kType = MMM::KeyframeType::FIRST;
    else if (ui->averageKeyframeRadioButton->isChecked()) kType = MMM::KeyframeType::AVERAGE;
    else kType = MMM::KeyframeType::LAST;
    std::set<std::string> segmentNames;

    for(int i = 0; i < ui->treeWidget->topLevelItemCount(); i++)
    {
       QTreeWidgetItem *item = ui->treeWidget->topLevelItem(i);
       if (item->checkState(0) == Qt::Checked) segmentNames.insert(item->text(1).toStdString());
    }

    if (segmentNames.size() == 0) {
        QMessageBox* msgBox = new QMessageBox(this);
        msgBox->setText("Choose some segments");
        msgBox->exec();
        return nullptr;
    }

    return MMM::PrincipalComponentAnalysisMotionSegmenterConfigurationPtr(new MMM::PrincipalComponentAnalysisMotionSegmenterConfiguration(alpha, reducedDimensions, kType, segmentNames));
}

void PrincipalComponentAnalysisMotionSegmenterWidget::setSegmentationType(MMM::SegmentationType &type) {
    this->segType = type;

    updateTreeWidget();
}

void PrincipalComponentAnalysisMotionSegmenterWidget::setMotion(MMM::MotionPtr motion) {
    this->motion = motion;

    if (motion) updateTreeWidget();
}

void PrincipalComponentAnalysisMotionSegmenterWidget::updateTreeWidget() {
    ui->treeWidget->clear();
    ui->treeWidget->setColumnCount(2);

    std::vector<std::string> segmentNames;
    switch (segType) {
    case MMM::SegmentationType::MMM:
        {
            ui->treeWidget->setHeaderLabels({QString("Use"), QString("Joint name")});
            MMM::KinematicSensorPtr kinematicSensor = motion->getSensorByType<MMM::KinematicSensor>(MMM::KinematicSensor::TYPE); // TODO!! More than one Kinematic Sensor
            if (!kinematicSensor) return;
            for (const std::string &jointName : kinematicSensor->getJointNames()) segmentNames.push_back(jointName);
        }
        break;
    case MMM::SegmentationType::MMM_MARKER:
        {
            ui->treeWidget->setHeaderLabels({QString("Use"), QString("MMM marker")});
            std::vector<VirtualRobot::SensorPtr> sensors = MMM::SimoxTools::buildModel(motion->getModel())->getSensors();
            for (auto sensor : sensors) {
                if (!sensor->getName().empty()) segmentNames.push_back(sensor->getName());
            }
        }
        break;
    case MMM::SegmentationType::C3D_MARKER:
        {
            ui->treeWidget->setHeaderLabels({QString("Use"), QString("C3D marker")});
            MMM::MoCapMarkerSensorPtr mocapSensor = motion->getSensorByType<MMM::MoCapMarkerSensor>(MMM::MoCapMarkerSensor::TYPE);
            if (!mocapSensor) return;
            MMM::MoCapMarkerSensorMeasurementPtr measurement = mocapSensor->getDerivedMeasurement(mocapSensor->getMinTimestep());
            for (const auto &labeledMarker : measurement->getLabeledMarker()) segmentNames.push_back(labeledMarker.first);
        }
        break;
    default:
        MMM_ERROR << "SegmentationType not initialized" << std::endl;
        break;
    }

    QList<QTreeWidgetItem*> items;
    for (const auto &name : segmentNames) {
        QTreeWidgetItem* item = new QTreeWidgetItem();
        item->setCheckState(0, Qt::Unchecked);
        item->setFlags(item->flags() | Qt::ItemIsUserCheckable);
        item->setText(1, QString::fromStdString(name));
        items.append(item);
    }

    ui->treeWidget->insertTopLevelItems(0, items);
    ui->treeWidget->expandAll();
    ui->treeWidget->resizeColumnToContents(0);
    ui->treeWidget->collapseAll();
}
